I've added a draft of the full functionality with U = U(y) (periodic in y). One question: when evaluating CalcN! there's an assignment Uy = real.(ifft(im * grid.l .* fft(params.U))). This seems to be required as there's no 1D fft plan in the grid that could be used with mul!. My questions:

1. Is this the best way to do this derivative?
2. An identical value is calculated and assigned each timestep. As it's only in the y direction it shouldn't take too much memory, however, is this step (in particular, assigning GPU memory) likely to slow down the timestepping?
3. We could have Uy saved in params but this seems unnecessary. Do you agree?

@mncrowe is it OK if I push few edits?

I think I see a minor problem. When U is just a number, don't we need a term - U ∂η/∂x? At the moment this term is not included.

I think I'll try to refactor the code to mimic the MultiLayerQG where the background PV gradients include the ∂η/∂x and ∂η/∂y terms and the ∂²U/∂y² (if U is not a constant).

@mncrowe I think I did everything! Have a look and let me know if you find any issues running the code or if the code is not behaving as you expect? Otherwise I'm happy to merge this!

I would just appreciate your look on the code. In principle I can approve the PR but it'd be nice if somebody else has a look because I made few changes.

Thanks for the very thorough reviewing and rewrite. Is there anything outstanding for me to do?

Does it work as expected?

A flow at rest advected by a background imposed U=0.5 over some topography.

using GeophysicalFlows, GLMakie, Printf

    dev = CPU()
      n = 256
stepper = "FilteredRK4"
     dt = 0.02
 nsteps = 4000
 nsubs  = 5
      L = 2π

σx, σy = 0.4, 0.8
topographicPV(x, y) = 3exp(-(x - 1)^2 / 2σx^2 - (y - 1)^2 / 2σy^2) - 2exp(- (x + 1)^2 / 2σx^2 - (y + 1)^2 / 2σy^2)

prob = SingleLayerQG.Problem(dev; nx=n, Lx=L, eta=topographicPV, U=0.5, dt, stepper, aliased_fraction=0)
sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

x,  y  = grid.x,  grid.y
Lx, Ly = grid.Lx, grid.Ly

q = Observable(Array(vars.q))

fig = Figure(size=(380, 300))

axis_kwargs = (xlabel = "x", ylabel = "y", aspect = 1,
               limits = ((-Lx/2, Lx/2), (-Ly/2, Ly/2)))

title_q = Observable("initial vorticity ∂v/∂x-∂u/∂y")
axq = Axis(fig[1, 1]; title = title_q, axis_kwargs...)

hm = heatmap!(axq, x, y, q, colormap = :balance, colorrange = (-5, 5))

Colorbar(fig[1, 2], hm)

contour!(axq, x, y, params.eta;
         levels = collect(0.5:1:2.5), linewidth = 2, color = (:black, 0.5))
contour!(axq, x, y, params.eta;
         levels = collect(-2.5:1:-0.5), linewidth = 2, color = (:grey, 0.7), linestyle = :dash)

title_q[] = "vorticity, t=" * @sprintf("%.2f", clock.t)

record(fig, "singlelayerqg_topography_U0.mp4", 0:round(Int, nsteps/nsubs), framerate = 40) do j
  q[] = irfft(sol, grid.nx)
  title_q[] = "vorticity, t=" * @sprintf("%.2f", clock.t)
  stepforward!(prob, nsubs)
end

Does it work as expected?

All my scripts are producing the same results as before, and match my earlier results using Dedalus.

Nice! I’m merging then.